/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "bsr_common.h"
#include "parbsr_blockjacobiilu_preconditioner.h"

#ifdef USE_OPENMP
#include <omp.h>
#endif

namespace YHAMG
{

ParBSR_BlockJacobiILU_Preconditioner::ParBSR_BlockJacobiILU_Preconditioner()
	: num_threads(0),
	D_LU(0)
{
}

ParBSR_BlockJacobiILU_Preconditioner::~ParBSR_BlockJacobiILU_Preconditioner()
{
	if (D_LU) delete[] D_LU;
}

void ParBSR_BlockJacobiILU_Preconditioner::Free()
{
	if (D_LU) delete[] D_LU;
	num_threads = 0;
	D_LU = 0;
	L.Free();
	U.Free();
}

static inline void Swap(int& a, int& b)
{
	int temp = a;
	a = b;
	b = temp;
}

#define CMP(a, b) ((a) < (b))

static void Sort(int* a, int left, int right)
{
	if (left >= right) return;

	Swap(a[left], a[(left + right) / 2]);

	int last = left;
	for (int i = left + 1; i <= right; ++i)
		if (CMP(a[i], a[left]))
			Swap(a[++last], a[i]);

	Swap(a[left], a[last]);

	Sort(a, left, last - 1);
	Sort(a, last + 1, right);
}

void ParBSR_BlockJacobiILU_Preconditioner::Setup(const ParBSR_Matrix& A, int REUSE)
{
	int n = A.LocLoc.size[0];
	int block_size = A.LocLoc.block_size;
	int block_nnz = block_size * block_size;
	int* Ap = A.LocLoc.rowptr;
	int* Ai = A.LocLoc.colind;
	double* Av = A.LocLoc.values;

	if (REUSE == 0)
	{
		Free();

		D_LU = new double[n * block_nnz];

		num_threads = 1;
#ifdef USE_OPENMP
		num_threads = omp_get_max_threads();
#endif

		int* w = new int[n];
		double* x = new double[n * block_nnz];

		for (int i = 0; i < n; ++i)
			w[i] = -1;

		int* Lp = new int[n + 1];
		int* Up = new int[n + 1];

		Lp[0] = 0;
		Up[0] = 0;
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
		for (int t = 0; t < num_threads; ++t)
		{
			int begin = t * n / num_threads;
			int end = (t + 1) * n / num_threads;

			for (int i = begin; i < end; ++i)
			{
				int cnt_l = 0;
				int cnt_u = 0;

				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
				{
					int jcol = Ai[j];

					if (jcol < begin || jcol >= end) continue;

					if (jcol < i)
						++cnt_l;
					else if (jcol > i)
						++cnt_u;
				}

				Lp[i + 1] = cnt_l;
				Up[i + 1] = cnt_u;
			}
		}

		for (int i = 0; i < n;++i)
			Lp[i + 1] += Lp[i];
		for (int i = 0; i < n;++i)
			Up[i + 1] += Up[i];

		int* Li = new int[Lp[n]];
		double* Lv = new double[Lp[n] * block_nnz];
		int* Ui = new int[Up[n]];
		double* Uv = new double[Up[n] * block_nnz];

#ifdef USE_OPENMP
#pragma omp parallel for
#endif
		for (int t = 0; t < num_threads; ++t)
		{
			int begin = t * n / num_threads;
			int end = (t + 1) * n / num_threads;

			for (int i = begin; i < end; ++i)
			{
				w[i] = i;
				BSRBlockFill(block_size, 0.0, x + i * block_nnz);

				for (int j = Ap[i], k = Lp[i], r = Up[i]; j < Ap[i + 1]; ++j)
				{
					int jcol = Ai[j];

					if (jcol < begin || jcol >= end) continue;

					w[jcol] = i;
					BSRBlockCopy(block_size, Av + j * block_nnz, x + jcol * block_nnz);

					if (jcol < i)
						Li[k++] = jcol;
					else if (jcol > i)
						Ui[r++] = jcol;
				}

				Sort(Li, Lp[i], Lp[i + 1] - 1);

				for (int k = Lp[i]; k < Lp[i + 1]; ++k)
				{
					int kcol = Li[k];

					BSRBlockMatLUSolve(block_size, D_LU + kcol * block_nnz, x + kcol * block_nnz);

					for (int j = Up[kcol]; j < Up[kcol + 1]; ++j)
					{
						int jcol = Ui[j];

						if (w[jcol] == i)
							BSRBlockMultSub(block_size, x + kcol * block_nnz, Uv + j * block_nnz, x + jcol * block_nnz);
					}
				}
				
				BSRBlockLUFactorize(block_size, x + i * block_nnz);
				BSRBlockCopy(block_size, x + i * block_nnz, D_LU + i * block_nnz);

				for (int j = Lp[i]; j < Lp[i + 1]; ++j)
					BSRBlockCopy(block_size, x + Li[j] * block_nnz, Lv + j * block_nnz);
				for (int j = Up[i]; j < Up[i + 1]; ++j)
					BSRBlockCopy(block_size, x + Ui[j] * block_nnz, Uv + j * block_nnz);
			}
		}

		L.size[0] = n;
		L.size[1] = n;
		L.block_size = block_size;
		L.rowptr = Lp;
		L.colind = Li;
		L.values = Lv;
		U.size[0] = n;
		U.size[1] = n;
		U.block_size = block_size;
		U.rowptr = Up;
		U.colind = Ui;
		U.values = Uv;

		delete[] w;
		delete[] x;
	}
	else
	{
		int* Lp = L.rowptr;
		int* Li = L.colind;
		double* Lv = L.values;
		int* Up = U.rowptr;
		int* Ui = U.colind;
		double* Uv = U.values;

		int* w = new int[n];
		double* x = new double[n * block_nnz];

		for (int i = 0; i < n; ++i)
			w[i] = -1;

#ifdef USE_OPENMP
#pragma omp parallel for
#endif
		for (int t = 0; t < num_threads; ++t)
		{
			int begin = t * n / num_threads;
			int end = (t + 1) * n / num_threads;

			for (int i = begin; i < end; ++i)
			{
				w[i] = i;
				x[i] = 0.0;

				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
				{
					int jcol = Ai[j];
					if (jcol >= begin && jcol < end)
					{
						w[jcol] = i;
						BSRBlockCopy(block_size, Av + j * block_nnz, x + jcol * block_nnz);
					}
				}

				for (int k = Lp[i]; k < Lp[i + 1]; ++k)
				{
					int kcol = Li[k];

					BSRBlockMatLUSolve(block_size, D_LU + kcol * block_nnz, x + kcol * block_nnz);

					for (int j = Up[kcol]; j < Up[kcol + 1]; ++j)
					{
						int jcol = Ui[j];

						if (w[jcol] == i)
							BSRBlockMultSub(block_size, x + kcol * block_nnz, Uv + j * block_nnz, x + jcol * block_nnz);
					}
				}

				BSRBlockLUFactorize(block_size, x + i * block_nnz);
				BSRBlockCopy(block_size, x + i * block_nnz, D_LU + i * block_nnz);

				for (int j = Lp[i]; j < Lp[i + 1]; ++j)
					BSRBlockCopy(block_size, x + Li[j] * block_nnz, Lv + j * block_nnz);
				for (int j = Up[i]; j < Up[i + 1]; ++j)
					BSRBlockCopy(block_size, x + Ui[j] * block_nnz, Uv + j * block_nnz);
			}
		}

		delete[] w;
		delete[] x;
	}
}

int ParBSR_BlockJacobiILU_Preconditioner::InSize() const
{
	return L.size[0] * L.block_size;
}

int ParBSR_BlockJacobiILU_Preconditioner::OutSize() const
{
	return L.size[0] * L.block_size;
}

void ParBSR_BlockJacobiILU_Preconditioner::Apply(const Par_Vector& b, const Par_Vector& x) const
{
	int n = L.size[0];
	int block_size = L.block_size;
	int block_nnz = block_size * block_size;
	int* Lp = L.rowptr;
	int* Li = L.colind;
	double* Lv = L.values;
	int* Up = U.rowptr;
	int* Ui = U.colind;
	double* Uv = U.values;
	double* bv = b.Loc.values;
	double* xv = x.Loc.values;

	if (n == 0) return;

	for (int i = 0; i < n; ++i)
		VecBlockCopy(block_size, bv + i * block_size, xv + i * block_size);

#ifdef USE_OPENMP
#pragma omp parallel for
#endif
	for (int t = 0; t < num_threads; ++t)
	{
		int begin = t * n / num_threads;
		int end = (t + 1) * n / num_threads;

		for (int i = begin; i < end; ++i)
			for (int j = Lp[i]; j < Lp[i + 1]; ++j)
				BSRBlockMatVecSub(block_size, Lv + j * block_nnz, xv + Li[j] * block_size, xv + i * block_size);

		for (int i = end - 1; i >= begin; --i)
		{
			for (int j = Up[i + 1] - 1; j >= Up[i]; --j)
				BSRBlockMatVecSub(block_size, Uv + j * block_nnz, xv + Ui[j] * block_size, xv + i * block_size);
			BSRBlockLUVecSolve(block_size, D_LU + i * block_nnz, xv + i * block_size);
		}
	}
}

}